%% Generate artificial data (Nh-CES).
% All matrices are in (I, N, T) order:
% I: the number of goods, N: the number of people, T: the number of time.
% Functions:
% Inputs: I_vec, pvec - exogenous wage and price (I, N, T)
% Outputs: u_vec - indirect utility, B_vec - uncompensated budget shares
% [u_vec, B_vec] = Gen_NhCES(I_vec, pvec, paras)

 
function [u_vec, B_vec] = Gen_NhCES(I_vec, pvec, paras)
% Nonhomothetic Preference: Given W and P, solve for u
% Initial value used to calculate the indirect utility of Comin
paras.uvec_init = I_vec / 10;

% Calibration of Omega to ensure budget shares of each good are the same for the median household
medianw = median(I_vec(1, :, 1));
omega_tmp = 1 ./ (medianw .^ ([0.2; 1; 1.65] * (1 - paras.sig)));
omega_init = omega_tmp ./ sum(omega_tmp);
omega_init2 = omega_init(1:end-1,:);
options = optimoptions('fsolve', 'Algorithm', 'trust-region', 'TolFun', 1e-7);
omega_tmp2 = fsolve(@(x) CalOmega(x, medianw, paras), omega_init2, options);
paras.omega = [omega_tmp2; 1 - sum(omega_tmp2)];
T = paras.T;
N = paras.N;
I = paras.I;
Ns = 50;
while mod(N, Ns) ~= 0
    Ns = Ns - 1;
end

v_vec = zeros(1,N,T);
uvec_init = paras.uvec_init;
options = optimoptions('fsolve', 'Algorithm', 'trust-region', 'TolFun', 1e-7,'Display','off');

% Loop through time and find the indirect utility for each hh
for t = 1:T
    for ns = 1:Ns
        Index = (N / Ns) * (ns - 1) + 1:(N / Ns) * (ns);
        [v_vec(:, Index, t), fval] = fsolve(@(x) NHU(x, I_vec(:, Index, t), pvec(:, :, t), paras), uvec_init(:, Index), options);
    end
end
%%Obtain EV
[u_vec,B_vec]= CalMM(v_vec,pvec,paras);

function err = NHU(v_vec,I_vec,pvec,paras) 
    E = sum(paras.omega.*(((v_vec.^paras.eps_vec).*pvec).^(1-paras.sig))).^(1./(1-paras.sig));
    err = abs(E-I_vec);
end

function [u_vec,B_vec]= CalMM(v_vec,pvec,paras)
    paras.sig =paras.sig;
    pvec_b = pvec(:,:,1);
    u_vec = sum(paras.omega.*(((v_vec.^paras.eps_vec).*pvec_b).^(1-paras.sig))).^(1./(1-paras.sig));
    tmp = paras.omega.*(((v_vec.^paras.eps_vec).*pvec).^(1-paras.sig));
    B_vec = tmp./sum(tmp);
end

function  err = CalOmega(omega_init,I_vec,paras) 
    % helper function for solve for Omega 
    options = optimoptions('fsolve', 'Algorithm', 'trust-region', 'TolFun', 1e-5, 'Display', 'off');
    [ v_vec ]  =  fsolve( @(x) NHU(x,I_vec,omega_init,paras),I_vec/10,options);  

    omega = [omega_init;1-sum(omega_init)]; 
    tmp = omega.*(((v_vec.^paras.eps_vec)).^(1-paras.sig));
    B_vec = tmp./sum(tmp);

    err = (abs(B_vec(1:end,:)-1/paras.I));
    
    function err = NHU(v_vec,I_vec,omega_init,paras) 
    omega = [omega_init;1-sum(omega_init)];
    E = sum(omega.*(((v_vec.^paras.eps_vec)).^(1-paras.sig))).^(1./(1-paras.sig));
    err = abs(E-I_vec);
end
end


end
